#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: main.py 
@time: 2022/04/04
@software: PyCharm 
description:
"""

import torch
from torch.utils.data import DataLoader

from models.model import TextClassification
from utils.data import IMDBDataset,to_map_style_dataset
from utils.utils import train_loop
from utils.parameters import Config


if __name__ == '__main__':
    config=Config()

    print(f"trian model with {config.device}")
    #模型，损失函数，优化器
    model=TextClassification(config)
    model=model.to(config.device)
    loss=torch.nn.CrossEntropyLoss()
    optimizer=torch.optim.SGD(model.parameters(),lr=config.lr)

    #数据集
    training_data=IMDBDataset(config)
    testing_data=IMDBDataset(config)

    train_dataloader=DataLoader(training_data,batch_size=config.batch_size,
                                 shuffle=True)
    test_dataloader = DataLoader(testing_data, batch_size=32,
                                   shuffle=True)

    #加载模型，进行训练
    train_loop(train_dataloader, model, loss, optimizer, config.save_path, config.save_dir, "", config.epochs)

    #模型预测
    # load_path='./model/2022-04-14/21-28-36.pt'
    # predict_loop(model,test_dataloader,load_path)

    print("done")